#!/usr/bin/env python
# vim: set fileencoding=UTF-8 :

# ------------------------------------------------------------
# ::                        WIGA-SOFT                       ::
# ------------------------------------------------------------
# loader.py
# Module for XLSX file parsing and inserting data to DB
# Created by: Vycas
# Version: 1.0
# ------------------------------------------------------------

from xml.dom.minidom import parse
from zipfile import ZipFile
from datetime import date
from exceptions import AttributeError
import shutil
from pao.models import Budget, Project, Employee, Department, Task
from django.contrib.auth.models import User
from django.contrib.auth.models import Group
from django.core.exceptions import ObjectDoesNotExist

class XLSXParser:
    """
    This class contains functions required to parse XLSX file.
    XLSX file have to be already extracted, so these functions
    work with XML files.

    readSharedStrings must be called before readSheet.
    """

    def readWorkbook(self, filename, sheetNames):
        """
        Reads workbook.xml file which reveals in which files are required sheets.
        Input: filename - Exact workbook.xml location.
               sheetNames - Dictionary where key is the sheet name in XLSX file
                            and value is the name of the returned sheet disctionary key.
        """

        sheets = {}
        dom = parse(filename)
        for sheet in dom.getElementsByTagName('sheet'):
            name = sheet.getAttribute('name')
            if name in sheetNames.keys():
                sheets[sheetNames[name]] = 'sheet' + sheet.getAttribute('r:id')[3] + '.xml'
        return sheets

    def readSharedStrings(self, filename):
        """
        Reads shared XLSX strings into list and returns it.
        Input: filename - XML file containing shared strings.
        Exceptions: IOError - In case when XML file doesn't exist.
        """
        
        self.sharedStrings = []
        dom = parse(filename)
        for node in dom.firstChild.childNodes:
            try:
                self.sharedStrings.append(node.firstChild.firstChild.nodeValue)
            except AttributeError:
                self.sharedStrings.append('')

    def readSheet(self, filename, columns):
        """
        Reads sheet from XML file and return a table (List of dictionaries).
        Input: filename - XML file containing tasks sheet.
               columns - Dictionary where key is column name in XLSX file and
                         value is the name of returned dictionary key.
        Exceptions: IOError - In case when XML file doesn't exist.
        """
 
        table = []
        dom = parse(filename)
        for row in dom.getElementsByTagName('row'):
            # Skip header
            if row.getAttribute('r') == '1':
                continue
            record = {}
            blank = True
            for c in columns.values():
                record[c] = None
            for col in row.childNodes:
                colName = col.getAttribute('r')[0:1]
                if not columns.has_key(colName):
                    continue
                try:
                    value = col.firstChild.firstChild.nodeValue
                    blank = False
                except AttributeError:
                    value = None
                if col.getAttribute('t') == 's':
                    record[columns[colName]] = self.sharedStrings[int(value)]
                else:
                    record[columns[colName]] = value
            if not blank:
                table.append(record)
        return table


class DataLoader:
    def __init__(self, xlsxFile):
        """
        Initializes DataLoader with input XLSX file. It can be string or file object.
        """

        self.xlsxFile = xlsxFile

    def __extract(self):
        """
        Extracts XLSX file to output directory.
        """

        zip = ZipFile(self.xlsxFile)
        zip.extractall(self.EXTRACT_DIR)
        zip.close()

    def __del__(self):
        """
        Removes temporary output directory, after reading XLSX file.
        """

        shutil.rmtree(self.EXTRACT_DIR)

    def __createGroups(self):
        """
        Creates user groups.
        """
        
        self.groupEmployee = Group.objects.get_or_create(name='Employee')[0]
        self.groupProjectManager = Group.objects.get_or_create(name='ProjectManager')[0]
        self.groupDepartmentManager = Group.objects.get_or_create(name='DepartmentManager')[0]
        self.groupEnterprise = Group.objects.get_or_create(name='Enterprise')[0]

    def __toAscii(self, word):
        """
        Removes lithuaninian chars.
        """
        word = word.replace(u'ą', 'a').replace(u'Ą', 'A')
        word = word.replace(u'č', 'c').replace(u'Č', 'C')
        word = word.replace(u'ę', 'e').replace(u'Ę', 'E')
        word = word.replace(u'ė', 'e').replace(u'Ė', 'E')
        word = word.replace(u'į', 'i').replace(u'Į', 'I')
        word = word.replace(u'š', 's').replace(u'Š', 'S')
        word = word.replace(u'ų', 'u').replace(u'Ų', 'U')
        word = word.replace(u'ū', 'u').replace(u'Ū', 'U')
        word = word.replace(u'ž', 'z').replace(u'Ž', 'Z')
        return word

    def updateDatabase(self):
        """
        Extracts XLSX file to temporary output directory, parses extracted
        XML files and updates database by inserting new records to it.
        """

        self.__extract()
        self.__createGroups()

        xlsx = XLSXParser()

        sharedStringsFilename = self.EXTRACT_DIR + 'xl/sharedStrings.xml'
        workbookFilename = self.EXTRACT_DIR + 'xl/workbook.xml'
        worksheetsPath = self.EXTRACT_DIR + 'xl/worksheets/'

        xlsx.readSharedStrings(sharedStringsFilename)
        sheets = xlsx.readWorkbook(workbookFilename, self.sheetNames)

        taskTable = xlsx.readSheet(worksheetsPath + sheets['Task'], self.taskColumns)
        projectTable = xlsx.readSheet(worksheetsPath + sheets['Project'], self.projectColumns)
        employeeTable = xlsx.readSheet(worksheetsPath + sheets['Employee'], self.employeeColumns)
        departmentTable = xlsx.readSheet(worksheetsPath + sheets['Department'], self.departmentColumns)
        budgetTable = xlsx.readSheet(worksheetsPath + sheets['Budget'], self.budgetColumns)
        
        for t in taskTable:
            t['employee'] = self.__toAscii(t['employee'])
        for p in projectTable:
            p['manager'] = self.__toAscii(p['manager'])
        for e in employeeTable:
            e['username'] = self.__toAscii(e['username'])
            e['email'] = self.__toAscii(e['email'])
        for d in departmentTable:
            d['manager'] = self.__toAscii(d['manager'])

        missingDepartments = {}
        users = User.objects.all()
        for employee in employeeTable:
            user = users.filter(username=employee['username'])
            if user.count() == 0:
                newUser = User(username=employee['username'], email=employee['email'])
                newUser.is_active = False
                newUser.is_staff = False
                newUser.save()
                newUser.groups.add(self.groupEmployee)
                newUser.save()
                newEmployee = Employee(id=newUser)
                newEmployee.save()
                missingDepartments[newEmployee] = employee['department']

        departments = Department.objects.all()
        for department in departmentTable:
            dep = departments.filter(id=department['department'])
            if dep.count() == 0:
                employee = User.objects.get(username=department['manager']).employee
                employee.id.groups.add(self.groupDepartmentManager)
                employee.save()
                newDepartment = Department(id=department['department'], name=department['name'], manager=employee)
                newDepartment.save()

        employees = missingDepartments.keys()
        for employee in employees:
            employee.department = Department.objects.get(id=missingDepartments[employee])
            employee.save()

        projects = Project.objects.all()
        users = User.objects.all()
        for project in projectTable:
            proj = projects.filter(id=project['project'])
            if proj.count() == 0:
                # If project is 'NEIVYKES' (there is no responsible manager) - database gets NULL
                try:
                    employee = users.get(username=project['manager']).employee
                    employee.id.groups.add(self.groupProjectManager)
                    employee.save()
                except ObjectDoesNotExist:
                    employee = None
                newProject = Project(id=project['project'], name=project['name'], manager=employee)
                newProject.save()

        projects = Project.objects.all()
        for budget in budgetTable:
            yearmonth = date(int(budget['year']), int(budget['month']), 1)
            budg = Budget.objects.filter(project=budget['project'], date=yearmonth)
            if budg.count() == 0:
                project = projects.get(id=budget['project'])
                workhours = int(budget['workhours'])
                newBudget = Budget(date=yearmonth, project=project, workhours=workhours)
                newBudget.save()

        for task in taskTable:
            yearmonth = date(int(task['year']), int(task['month']), 1)
            tsk = Task.objects.filter(project=task['project'], date=yearmonth, employee__id__username=task['employee'])
            if tsk.count() == 0:
                project = projects.get(id=task['project'])
                employee = users.get(username=task['employee']).employee
                workhours = int(task['workhours'])
                name = task['name']
                newTask = Task(date=yearmonth, project=project, employee=employee, workhours=workhours, name=name)
                newTask.save()

    EXTRACT_DIR = 'extracted/'
    

    sheetNames = {u'Užduotys': 'Task',
                  u'Projektai': 'Project',
                  u'Darbuotojai': 'Employee',
                  u'Skyriai': 'Department',
                  u'Biudžetai': 'Budget'}

    taskColumns = {'A': 'project',
                   'B': 'employee',
                   'C': 'year',
                   'D': 'month',
                   'E': 'workhours',
                   'F': 'name'}

    projectColumns = {'A': 'project',
                      'B': 'manager',
                      'C': 'name'}

    employeeColumns = {'A': 'username',
                       'B': 'department',
                       'C': 'email'}

    departmentColumns = {'A': 'department',
                         'B': 'manager',
                         'C': 'name'}

    budgetColumns = {'A': 'project',
                     'B': 'year',
                     'C': 'month',
                     'D': 'workhours'}
